"""
构建类型训练集，也就是输入一句话，询问是什么类型
之前的example007不符合要求（没”其他“类型，除了filter_mapping.json中的几种类型，应该都归类为“其他”）
需要添加其他类型，其他类型和正常的类型比例为1：1

"""

import json
import os
import re
from collections import OrderedDict
import random

from extract_attributes.main import get_allFile_from_folder, read_pdf_from_file
from pdf2excelv2.get_attributes import get_star_structure, replace_from_text, out2file
from pdf2excelv2.main import text2sentence


# 通过excel文件得到星型结构，参照003
def star_to_type_train(dic):
    tep_dics = []
    sett = set()
    for key, value in dic.items():
        text = value["描述"]

        # 使用正则表达式查找answer中的中文
        pattern = re.compile(r'[\u4e00-\u9fa5]+')  # 中文字符的Unicode编码范围是0x4e00到0x9fa5
        chinese_chars = pattern.findall(str(key))

        # 将中文字符合并成一个字符串
        chinese_str = ''.join(chinese_chars)

        print(chinese_str, '++++++++')

        # 读filter_mapping.json，获取到所有的键，只保留这些键
        with open('filter_mapping.json', 'r', encoding='utf-8') as f:
            data = json.load(f)
        keys = data.keys()
        # excel的property属性 不是filter_mappgin.json key值之一
        if chinese_str not in keys:
            continue

        tep_dic = {"ask": "根据文本内容:“" + str(text) + "”回答实体类型是什么", "answer": chinese_str}
        sett.add(chinese_str)
        tep_dics.append(tep_dic)
    lst = list({tuple(d.items()) for d in tep_dics})
    tep_dics = [dict(item) for item in lst]

    # tep_dics = list(OrderedDict.fromkeys(map(tuple, tep_dics)))
    print(sett)
    return tep_dics


def normal_type_train(folder_path=r'指南知识图谱三元组提取'):
    pdf_file_paths = (get_allFile_from_folder(folder_path, file_type="pdf"))
    excel_file_paths = (get_allFile_from_folder(folder_path, file_type="xlsx"))
    results = []
    for pdf_file_path in pdf_file_paths:
        dirname = os.path.dirname(pdf_file_path)
        #  获取相同文件下的excel文件
        excel_file_path = [i for i in excel_file_paths if os.path.dirname(i) == dirname][0]
        orgin_dic = get_star_structure(excel_file_path)
        replace_dic = replace_from_text(pdf_file_path, orgin_dic)

        results = results + star_to_type_train(replace_dic)+star_to_type_train(orgin_dic)

    return results


"""
获得所有other类型的数据
1、将pdf分句
2、如果不存在results（filter_mapping.json中的几种类型的数据）中，就是其他

输入：文件夹名称
输出：其他类型的数据
"""


def other_type_train(results, folder_path=r'指南知识图谱三元组提取'):
    # 读所有的pdf文件名
    pdf_file_paths = (get_allFile_from_folder(folder_path, file_type="pdf"))
    others = []
    for pdf_file_path in pdf_file_paths:
        # 通过文件名得到pdf的内容
        pdf_text = read_pdf_from_file(pdf_file_path).replace("\n", "").encode('gb2312', 'ignore').decode('gb2312')
        # 分句
        pdf_sentences = text2sentence(pdf_text)
        # 在result中不存在的认为是其他
        for pdf_sentence in pdf_sentences:
            pdf_sentence = "根据文本内容:“" + str(pdf_sentence) + "”回答实体类型是什么"
            # 利用列表表达式获得dic中的ask
            asks = [d['ask'] for d in results]
            if pdf_sentence not in asks:
                tep_dic = {"ask": pdf_sentence, "answer": "其他"}
                others.append(tep_dic)
    # 取全量的“其他”样本，其他样本太多，导致泛化性太差，所以需要随机取部分
    return random.sample(others, round(len(others) * 0.4))


"""
正常的和其他的和
"""
def get_all_train(folder_path=r'指南知识图谱三元组提取'):
    normal = normal_type_train(folder_path=folder_path)
    other = other_type_train(normal, folder_path=folder_path)
    res =normal + other
    random.shuffle(res)
    out2file("type_traindata0001.json", res)


"""
构建实体类型训练集
"""
def construct_type_traindata(traindata_path):
    type_list = []
    with open(r'type_traindata003.json', "a+", encoding="utf-8") as f1:
        with open(traindata_path, "r", encoding='utf-8') as f:
            data = json.load(f)
        for dict in data:
            input_str = re.split(r'^(根据文本内容:)', dict['ask'])[2]
            # 去掉"
            instruction_str = re.split(r'(回答实体类型是什么)$', input_str)[0].strip('“')
            tep_dic = {"instruction": "你是一个实体类型识别模型，请根据下面文本回答实体类型是什么，文本：", "input": str(instruction_str), "output": str(dict['answer']) }
            type_list.append(tep_dic)
        f1.write(json.dumps(type_list, ensure_ascii=False, indent=4))

if __name__ == '__main__':
    get_all_train(folder_path=r'../指南知识图谱三元组提取')
    # construct_type_traindata('./type_traindata.json')